import numpy as np
import matplotlib.pyplot as plt
import seaborn as sns

# Set global font to Times New Roman
plt.rcParams['font.family'] = 'Times New Roman'
plt.rcParams['font.size'] = 12


def sinusoidal_position_embeddings(times, dim):
    half_dim = dim // 2
    embeddings = np.zeros((len(times), dim))

    for pos in range(len(times)):
        for i in range(0, dim, 2):
            omega = 1.0 / (10000 ** ((2 * i) / dim))
            embeddings[pos, i] = np.sin(times[pos] * omega)
            embeddings[pos, i + 1] = np.cos(times[pos] * omega)

    return embeddings


# Generate time steps and embeddings
dim = 64
times = np.linspace(0, 100, 50)
embeddings = sinusoidal_position_embeddings(times, dim)

# 1. Heatmap visualization
plt.figure(figsize=(12, 6))
sns.heatmap(embeddings, cmap='viridis')
plt.title('Time Embedding Heatmap Visualization')
plt.xlabel('Embedding Dimension')
plt.ylabel('Time Step')
plt.show()

# 2. Time encoding curves for different dimensions
plt.figure(figsize=(12, 6))
for i in range(0, dim, 8):
    plt.plot(times, embeddings[:, i], label=f'Dimension {i}')
plt.title('Time Encoding Patterns')
plt.xlabel('Time Step')
plt.ylabel('Encoding Value')
plt.legend()
plt.grid(True)
plt.show()

# 3. 3D visualization
from mpl_toolkits.mplot3d import Axes3D

fig = plt.figure(figsize=(10, 8))
ax = fig.add_subplot(111, projection='3d')

dim1, dim2, dim3 = 0, 16, 32
ax.plot3D(embeddings[:, dim1],
          embeddings[:, dim2],
          embeddings[:, dim3],
          'blue')

ax.set_xlabel(f'Dimension {dim1}')
ax.set_ylabel(f'Dimension {dim2}')
ax.set_zlabel(f'Dimension {dim3}')
ax.set_title('Time Encoding 3D Trajectory')
plt.show()

# 4. Frequency characteristics visualization
plt.figure(figsize=(12, 6))
dense_times = np.linspace(0, 10, 1000)
dense_embeddings = sinusoidal_position_embeddings(dense_times, dim)

for i in range(0, 8, 2):
    plt.plot(dense_times, dense_embeddings[:, i],
             label=f'Dim {i} (Low Freq)')
    plt.plot(dense_times, dense_embeddings[:, -i - 1] if i > 0 else dense_embeddings[:, -1],
             '--', label=f'Dim {-i - 1 if i > 0 else -1} (High Freq)')

plt.title('Frequency Characteristics of Time Encoding')
plt.xlabel('Time Step')
plt.ylabel('Encoding Value')
plt.legend()
plt.grid(True)
plt.show()

# 5. Phase relationship visualization
plt.figure(figsize=(12, 6))
selected_dims = [0, 1, 16, 17, 32, 33]
time_point = 25

for i in range(0, len(selected_dims), 2):
    dim1, dim2 = selected_dims[i], selected_dims[i + 1]
    plt.scatter(embeddings[:, dim1], embeddings[:, dim2],
                label=f'Dims {dim1}-{dim2}', alpha=0.5)
    plt.scatter(embeddings[time_point, dim1], embeddings[time_point, dim2],
                c='red', s=100, marker='x')

plt.title('Phase Relationships in Time Encoding')
plt.xlabel('sin(ωt)')
plt.ylabel('cos(ωt)')
plt.legend()
plt.grid(True)
plt.axis('equal')
plt.show()